package com.xu.cloud.common.limit;

import com.xu.cloud.common.exception.ErrorCodeEnum;
import com.xu.cloud.common.api.ApiContext;
import com.xu.cloud.common.util.ApplicationContextUtil;
import com.xu.cloud.common.util.HttpServletUtil;
import com.xu.cloud.common.util.IpUtil;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aop.Pointcut;
import org.springframework.aop.framework.autoproxy.AbstractBeanFactoryAwareAdvisingPostProcessor;
import org.springframework.aop.support.ComposablePointcut;
import org.springframework.aop.support.DefaultPointcutAdvisor;
import org.springframework.aop.support.annotation.AnnotationMatchingPointcut;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.DefaultParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.util.Assert;

import java.lang.reflect.Method;
import java.time.ZoneOffset;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * 限流aop拦截器
 *
 * @author xuguan
 * @since 2024/9/24
 */
@Configuration
public class RateLimitAdvisingBeanPostProcessor extends AbstractBeanFactoryAwareAdvisingPostProcessor {

    public static final Logger log = LoggerFactory.getLogger(RateLimitAdvisingBeanPostProcessor.class);

    @Override
    public void setBeanFactory(BeanFactory beanFactory) {
        super.setBeanFactory(beanFactory);
        final Pointcut cpc = new AnnotationMatchingPointcut(RateLimit.class, true);
        final Pointcut mpc = new AnnotationMatchingPointcut(null, RateLimit.class, true);
        final ComposablePointcut composablePointcut = new ComposablePointcut(cpc).union(mpc);
        this.advisor = new DefaultPointcutAdvisor(composablePointcut, new RateLimitMethodInterceptor());
    }

    private static class RateLimitMethodInterceptor implements MethodInterceptor {
        @Override
        public Object invoke(MethodInvocation methodInvocation) throws Throwable {
            final RateLimit methodRateLimit = AnnotationUtils.findAnnotation(methodInvocation.getMethod(), RateLimit.class);
            final RateLimit typeRateLimit = AnnotationUtils.findAnnotation(methodInvocation.getThis().getClass(), RateLimit.class);
            // 方法上注解优先级高于类上注解
            RateLimit rateLimit = Objects.nonNull(methodRateLimit)
                    ? methodRateLimit
                    : typeRateLimit;

            if (Objects.isNull(rateLimit)) {
                return methodInvocation.proceed();
            }

            final String rateLimitKey = getRateLimitKey(rateLimit, methodInvocation);
            final int limitCount = rateLimit.limitCount();
            final int limitDuration = rateLimit.limitDuration();
            final TimeUnit timeUnit = rateLimit.timeUnit();
            final boolean limitIfReqFail = rateLimit.limitIfReqFail();

            if (RateLimit.NOT_LIMIT == limitCount || RateLimit.NOT_LIMIT == limitDuration) {
                return methodInvocation.proceed();
            }

            // 校验限流
            checkRateLimit(rateLimitKey, limitCount, limitDuration, timeUnit);

            final Object proceed;
            try {
                proceed = methodInvocation.proceed();

                // 添加限流
                addRateLimitKey(rateLimitKey, limitDuration, timeUnit);
                return proceed;
            } catch (Throwable e) {
                // 请求失败是否添加限流
                if (limitIfReqFail) {
                    addRateLimitKey(rateLimitKey, limitDuration, timeUnit);
                }
                throw e;
            }
        }

        private String getRateLimitKey(final RateLimit rateLimit, final MethodInvocation methodInvocation) {
            final String expression = rateLimit.expression();

            final Object target = methodInvocation.getThis();
            final Object[] args = methodInvocation.getArguments();
            final Method method = methodInvocation.getMethod();
            final ParameterNameDiscoverer discoverer = new DefaultParameterNameDiscoverer();
            final String[] parameterNames = discoverer.getParameterNames(method);

            final StandardEvaluationContext context = new StandardEvaluationContext();
            for (int i = 0; Objects.nonNull(parameterNames) && i < parameterNames.length; i++) {
                context.setVariable(parameterNames[i], args[i]);
            }

            final SpelExpressionParser parser = new SpelExpressionParser();
            String value = parser.parseExpression(expression).getValue(context, String.class);
            if (Objects.equals(parser.parseExpression(RateLimit.CLIENT_IP).getValue(context, String.class), value)) {
                value = IpUtil.getClientIp(HttpServletUtil.currentRequest());
            }

            Assert.notNull(value, "rate limit expression value must not be null");

            final String redisPrefix = rateLimit.prefix();
            final String canonicalName = target.getClass().getCanonicalName();
            final String methodName = method.getName();
            return redisPrefix + canonicalName + ":" + methodName + ":" + value;
        }

        private void checkRateLimit(String rateLimitKey, int limitCount, int limitDuration, TimeUnit timeUnit) {
            Assert.isTrue(limitCount > 0, "rate limit count must greater than zero");
            Assert.isTrue(limitDuration > 0, "rate limit time must greater than zero");

            log.info("限流配置生效, 缓存key={}, 限流次数={}, 限流时间段={}({})",
                    rateLimitKey, limitCount, limitDuration, timeUnit.name().toLowerCase());

            final long reqMilli = ApiContext.current().getReqTime().toInstant(ZoneOffset.of("+8")).toEpochMilli();
            final long durationMillis = timeUnit.toMillis(limitDuration);

            final StringRedisTemplate redisTemplate = ApplicationContextUtil.getBean(StringRedisTemplate.class);
            final ZSetOperations<String, String> zSetOperations = redisTemplate.opsForZSet();

            // 清空限流时间段之前的值
            zSetOperations.removeRangeByScore(rateLimitKey, 0, reqMilli - durationMillis);

            // 获取缓存中请求次数
            final Long count = zSetOperations.count(rateLimitKey, 0, reqMilli);
            Assert.notNull(count, "redis zSet count must not be null");
            log.info("限流时间段内请求次数为count={}", count);

            if (count >= limitCount) {
                throw ErrorCodeEnum.ERR_RATE_LIMIT.info().newException();
            }
        }

        private void addRateLimitKey(String rateLimitKey, int limitDuration, TimeUnit timeUnit) {
            final StringRedisTemplate redisTemplate = ApplicationContextUtil.getBean(StringRedisTemplate.class);
            final ZSetOperations<String, String> zSetOperations = redisTemplate.opsForZSet();
            final long reqMilli = ApiContext.current().getReqTime().toInstant(ZoneOffset.of("+8")).toEpochMilli();

            zSetOperations.add(rateLimitKey, String.valueOf(reqMilli), reqMilli);
            redisTemplate.expire(rateLimitKey, limitDuration, timeUnit);
        }
    }
}

